#include "TestState.h"
#include "DX11ErrorReport.h"
#include "DX11HelpfulFunctions.h"
#include "Timer.h"
#include "CameraManager.h"
#include "LightingManager.h"
#include "Vec3.h"
#include "Game.h"
#include "AntTweakBarManager.h"
#include "DX11TextureManager.h"
#include "DX11ShaderManager.h"
#include "Text.h"
#include "ConfigManager.h"
#include "Get3DMousePos.h"
#include "OptionsState.h"

// ray tracing includes

#include "Colour.h"
#include "Ray.h"
#include "Camera.h"
#include "RTLight.h"
#include "Sphere.h"
#include "Plane.h"
#include "Triangle.h"

#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <math.h>
#include <limits>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>

// ray tracing includes end

// compute shader includes

#include "ComputeShader.h"

// compute shader includes end


TestState::TestState()
{
	m_rotCrateAngle = 0;
	for (int i = 0 ; i != 15; i++)
	{
		axis[i].x = 0;
		axis[i].y = 0;
		axis[i].z = 0;
	}

	m_screenX = TheScreenManager::Instance()->GetWindowWidth();
	m_screenY = TheScreenManager::Instance()->GetWindowHeight();
	m_retKeyDown = false;
}

TestState::~TestState()
{
}

// function calculates final colour of the pixel
// use to implement shadows, reflections etc
Colour GetColour(const D3DXVECTOR3& point, const D3DXVECTOR3& dir, int objectIndex, 
				 std::vector<RTLight*>& lights, std::vector<Object*>& objects, float accuracy, float ambientLight)
{
	Colour objectColour = objects[objectIndex]->colour;
	D3DXVECTOR3 finalObjectNormal = objects[objectIndex]->GetNormalAt(point);

	Colour finalColour = objectColour.Scalar(ambientLight);

	int numOfLights = lights.size();
	for (int i = 0; i != numOfLights; i++)
	{
		D3DXVECTOR3 lightDirection = lights[i]->m_position - point;
		D3DXVec3Normalize(&lightDirection, &lightDirection);

		float cosAngle = D3DXVec3Dot(&finalObjectNormal, &lightDirection);

		if (cosAngle > 0.0f)
		{
			// test for shadows
			bool inShadow = false;

			D3DXVECTOR3 directionFromLight = lights[i]->m_position - point;
			D3DXVec3Normalize(&directionFromLight, &directionFromLight);

			float distFromLight = D3DXVec3LengthSq(&directionFromLight);

			D3DXVECTOR3 shadowRayDir = lights[i]->m_position - point;
			D3DXVec3Normalize(&shadowRayDir, &shadowRayDir);
			Ray shadowRay(point, shadowRayDir);

			std::vector<float> shadowIntersections;

			for (int j = 0; j != objects.size(); j++)
			{
				shadowIntersections.push_back(objects[j]->FindIntersection(shadowRay));
			}

			for (int k = 0; k != shadowIntersections.size(); k++)
			{
				if (shadowIntersections[k] > accuracy)
				{
					if (shadowIntersections[k] <= distFromLight)
					{
						inShadow = true;
					}
					break;
				}
			}

			if (!inShadow)
			{
				//finalColour = finalColour * lights[i]->m_colour.Scalar(cosAngle);
				finalColour += objectColour * lights[i]->m_colour.Scalar(cosAngle);
				
				float intersectedObjectGloss = objects[objectIndex]->colour.gloss;
				if (intersectedObjectGloss > 0.0f && intersectedObjectGloss < 1.0f)
				{
					float dot = D3DXVec3Dot(&finalObjectNormal, &dir);
					D3DXVECTOR3 scalar1 = finalObjectNormal * dot;
					D3DXVECTOR3 final1 = scalar1 + dir;

					D3DXVECTOR3 scalar2 = final1 * 2.0f;
					D3DXVECTOR3 final2 = scalar2 - dir;

					D3DXVECTOR3 reflectionDir;
					D3DXVec3Normalize(&reflectionDir, &final2);

					float specular = D3DXVec3Dot(&reflectionDir, &dir);
					if (specular > 0.0f)
					{
						specular = pow(specular, 800.0f);
						finalColour += lights[i]->m_colour.Scalar(specular * intersectedObjectGloss);
					}
				}
			}
		}
	}



	return finalColour.Clip();
}

float sphere1Radi = 1.0f;
float sphere2Radi = 0.5f;

D3DXVECTOR3 sphere1Pos(-1.75f, 0.25f, 0.0f);
D3DXVECTOR3 sphere2Pos(1.75f, -0.25f, 0.0f);
D3DXVECTOR3 plane1Pos(0.0f, 1.0f, 0.0f);

bool TestState::SaveImg(const std::string& filename)
{
	int k = m_screenX * m_screenY;

	Colour* data = new Colour[k];
	int currentPixel = 0;
	float aspectRatio = TheScreenManager::Instance()->GetAspectRatio();

	float ambientLight = 0.2f;
	float accuracy = 0.000001f;
	//double accuracy = 0.00000001;


	// setting up camera

	D3DXVECTOR3 X(1.0f, 0.0f, 0.0f);
	D3DXVECTOR3 Y(0.0f, 1.0f, 0.0f);
	D3DXVECTOR3 Z(0.0f, 0.0f, 1.0f);

	Camera cam;

	{
		//D3DXVECTOR3 camPos = TheCameraManager::Instance()->GetCameraPos();
		D3DXVECTOR3 camPos = D3DXVECTOR3(3.0f, 1.5f, -4.0f);
		//D3DXVECTOR3 lookAt = TheCameraManager::Instance()->GetLookAt();
		D3DXVECTOR3 lookAt = D3DXVECTOR3(0.0f, 0.0f, 0.0f);
		D3DXVECTOR3 diffBetween = camPos - lookAt;
		D3DXVECTOR3 camDirection = -diffBetween;
		D3DXVec3Normalize(&camDirection, &camDirection);

		D3DXVECTOR3 camRight;
		D3DXVECTOR3 camDown;


		D3DXVec3Cross(&camRight, &Y, &camDirection);
		D3DXVec3Normalize(&camRight, &camRight);
		D3DXVec3Cross(&camDown, &camRight, &camDirection);
		//D3DXVec3Normalize(&camDown, &camDown);
		Camera camInit(camPos, camDirection, camRight, camDown);
		cam = camInit;
	}

	Colour white(1.0f, 1.0f, 1.0f, 1.0f);
	Colour green(0.5f, 1.0f, 0.5f, 1.0f);
	green.gloss = 0.3f;
	Colour maroon(0.5f, 0.25f, 0.25f, 1.0f);
	Colour red(1.0f, 0.0f, 0.0f, 1.0f);
	Colour grey(0.5f, 0.5f, 0.5f, 1.0f);
	Colour black(0.0f, 0.0f, 0.0f, 1.0f);

	D3DXVECTOR3 lightPos(-7.0f, 10.0f, -10.0f);
	RTLight light1(lightPos, white);

	std::vector<RTLight*> lightSources;
	lightSources.push_back(&light1);

	// setting up camera end

	// scene objects
	Sphere sphere1(sphere1Pos, sphere1Radi, green);
	Sphere sphere2(sphere2Pos, sphere2Radi, maroon);
	Triangle tri(D3DXVECTOR3(2, -0.5, 0), D3DXVECTOR3(0, -1, 0), D3DXVECTOR3(0, -0.5, -2), green);
	Plane plane1(plane1Pos, -1.0f, red);
	
	std::vector<Object*> objects;
	objects.push_back(dynamic_cast<Object*>(&sphere1));
	objects.push_back(dynamic_cast<Object*>(&sphere2));
	objects.push_back(dynamic_cast<Object*>(&tri));
	objects.push_back(dynamic_cast<Object*>(&plane1));

	// scene obejcts end


	// write image colour start
	
	float xAmount = 0.0f;
	float yAmount = 0.0f;

	for (int x = 0; x < m_screenX; x++)
	{
		for (int y = 0; y != m_screenY; y++)
		{
			if (y == 479 && x == 0)
			{
				std::cout << "stop!!!" << std::endl;
			}

			currentPixel = y * m_screenX + x;

			// start with no anti-aliasing
			if (m_screenX > m_screenY) // can change for aspect ratio later
			{
				// wide image
				xAmount = ((x + 0.5f) / m_screenX) * aspectRatio - ((m_screenX - m_screenY) / (float)m_screenY / 2.0f);
				yAmount = ((m_screenY - y) + 0.5f) / m_screenY;
			}
			else if (m_screenY > m_screenX)
			{
				// tall image
				xAmount = (x + 0.5f) / m_screenX;
				yAmount = (((m_screenY - y) + 0.5f) / m_screenY) / aspectRatio - (((m_screenY - m_screenX) / (float)m_screenX) / 2.0f);
			}
			else
			{
				// square image
				xAmount = (x + 0.5f) / m_screenX;
				yAmount = ((m_screenY + y) + 0.5f) / m_screenY;
			}

			// ray calculations cast by the camera position
			D3DXVECTOR3 rayOrigin = cam.GetPosition();
			D3DXVECTOR3 rayDirection = cam.GetDirection();// + ((camRight * (xAmount - 0.5f)) + ((camDown * (yAmount - 0.5f))));
			rayDirection += (cam.GetRight() * (xAmount - 0.5f)) + (cam.GetDown() * (yAmount - 0.5f));
			D3DXVec3Normalize(&rayDirection, &rayDirection);

			// the ray used to test collisions with scene objects
			Ray camRay(rayOrigin, rayDirection);

			std::vector<float> intersections;
			// test if ray intersects with objects in scene
			for (int o = 0; o != objects.size(); o++)
			{
				intersections.push_back(objects[o]->FindIntersection(camRay));
			}

			int finalIntersectedObject = -1;
			if (intersections.size() == 1)
			{
				if (intersections[0] > 0.0f)
				{
					finalIntersectedObject = 0;
				}
			}
			else
			{
				// more than one intersection the ray encountered so find the one closest to the camera
				float minVal = 10000.0f; // culled distance
				int index = 0;
				for (int i = 0; i != intersections.size(); i++)
				{
					if (minVal > intersections[i] && intersections[i] != -1.0f && intersections[i] < 10000.0f)
					{
						minVal = intersections[i];
						index = i;
					}
				}

				if (minVal > 0.0f)
				{
					finalIntersectedObject = index;
				}
			}

			//std::cout << finalIntersectedObject;

			if (finalIntersectedObject == -1)
			{
				data[currentPixel].r = 0.0f;
				data[currentPixel].g = 0.0f;
				data[currentPixel].b = 0.0f;
			}
			else
			{
				// check if greater than accuracy, eliminates bad shadows
				if (intersections[finalIntersectedObject] > accuracy)
				{
					//if (finalIntersectedObject == 1)
					//{
					//	// check id
					//	// check intersection val
					//	std::cout << "lol" << std::endl;
					//}


					D3DXVECTOR3 intersectedPoint = rayOrigin + (rayDirection * intersections[finalIntersectedObject]);
					D3DXVECTOR3 intersectedRayDir = rayDirection;

					Colour intersectedColour = GetColour(intersectedPoint, intersectedRayDir, finalIntersectedObject, 
															lightSources, objects, accuracy, ambientLight);
					//intersectedColour = objects[finalIntersectedObject]->colour;
					data[currentPixel] = intersectedColour;
				}
			}
		}
	}

	// write file colour end

	FILE *f;
	int dpi = 72;
	int size = k * 4;
	int fileSize = 54 + size; // 54 because 14 + 40 chars for file+info headers

	float factor = 39.375f;
	int m =	static_cast<int>(factor);
	int ppm = dpi * m;
	
	unsigned char bmpfileheader[14] = {'B', 'M', 0, 0, 0, 0,
												0, 0, 0, 0,
												54, 0, 0, 0};

	unsigned char bmpinfoheader[40] = {40, 0, 0, 0,
										0, 0, 0, 0,
										0, 0, 0, 0,
										1, 0, 24, 0};

	bmpfileheader[ 2] = (unsigned char)(fileSize);
	bmpfileheader[ 3] = (unsigned char)(fileSize >> 8);
	bmpfileheader[ 4] = (unsigned char)(fileSize >> 16);
	bmpfileheader[ 5] = (unsigned char)(fileSize >> 24);
	
	bmpinfoheader[ 4] = (unsigned char)(m_screenX);
	bmpinfoheader[ 5] = (unsigned char)(m_screenX >> 8);
	bmpinfoheader[ 6] = (unsigned char)(m_screenX >> 16);
	bmpinfoheader[ 7] = (unsigned char)(m_screenX >> 24);
	
	bmpinfoheader[ 8] = (unsigned char)(m_screenY);
	bmpinfoheader[ 9] = (unsigned char)(m_screenY >> 8);
	bmpinfoheader[10] = (unsigned char)(m_screenY >> 16);
	bmpinfoheader[11] = (unsigned char)(m_screenY >> 24);
	
	bmpinfoheader[25] = (unsigned char)(ppm);
	bmpinfoheader[26] = (unsigned char)(ppm >> 8);
	bmpinfoheader[27] = (unsigned char)(ppm >> 16);
	bmpinfoheader[28] = (unsigned char)(ppm >> 24);
	
	bmpinfoheader[29] = (unsigned char)(ppm);
	bmpinfoheader[30] = (unsigned char)(ppm >> 8);
	bmpinfoheader[31] = (unsigned char)(ppm >> 16);
	bmpinfoheader[32] = (unsigned char)(ppm >> 24);

	fopen_s(&f, filename.c_str(), "wb");

	fwrite(bmpfileheader, 1, 14, f);
	fwrite(bmpinfoheader, 1, 40, f);

	for (int i = 0; i != k; i++) 
	{
		Colour col = data[i];
		
		float r = data[i].r * 255.0f;
		float g = data[i].g * 255.0f;
		float b = data[i].b * 255.0f;

		//int colour[3] = {1.0f, 0.0f, 1.0f};
		unsigned char colour[3] = {(int)floor(b), (int)floor(g), (int)floor(r)};
		fwrite(colour, 1, 3, f);
	}

	fclose(f);
	m_retKeyDown = false;

	delete[] data;

	std::cout << "finished" << std::endl;

	return true;
}

void TestState::Draw()
{
	DX11::ClearColour(0.5f, 0.5f, 0.5f, 1.0f);
	
	//TheGame::Instance()->DrawColourPickObjects();
	
	
	TheScreenManager::Instance()->GetDeviceContext()->OMSetRenderTargets(1, 
		TheScreenManager::Instance()->GetRenderTargetview(),
		TheScreenManager::Instance()->GetDepthStencilView());		
	DX11::ClearColour(0.0f, 0.0f, 0.0f, 1.0f);
	TheCameraManager::Instance()->Update();
	
	DX11::ClearColour(0.0f, 0.0f, 0.0f, 1.0f);

	TheCameraManager::Instance()->Draw();
	TheGame::Instance()->DrawGameObjects();

	// start ray trace tests
	
	static Text t;
	if (m_retKeyDown)
	{
		//SaveImg("TestImage.bmp");
		t.Update("Finished Rendering");
	}
	
	// end ray trace tests

	static ComputeShader cs;
	static bool loadcs = cs.Load("../../DX11Shaders/TestCS2.hlsl");

	//D3D11_FEATURE_DATA_THREADING data;
	//HRESULT hr = TheScreenManager::Instance()->GetDevice()->CheckFeatureSupport(D3D11_FEATURE_THREADING, &data, sizeof(D3D11_FEATURE_DATA_THREADING));
	
	D3DXMATRIX m_view;
	D3DXMATRIX m_proj;

	D3DXMatrixLookAtLH(&m_view, &D3DXVECTOR3(0.0f, 0.0f, 1.0f), 
		&D3DXVECTOR3(0.0f, 0.0f, 0.0f), 
		&D3DXVECTOR3(0.0f, 1.0f, 0.0f));
	
	D3DXMatrixOrthoLH(&m_proj, TheScreenManager::Instance()->GetWindowWidth(), 
		TheScreenManager::Instance()->GetWindowHeight(), 0, 100);

	TheScreenManager::Instance()->SetViewProjMatrix(m_view, m_proj);
	
	cs.Run();

	

	//// set texture
	////TheObjMeshManager::Instance()->ChangeTexture("title_quad",
	////	TheDX11TextureManager::Instance()->UseTexture(L"../../Assets/title_image.png"));
	//TheObjMeshManager::Instance()->ChangeTexture("title_quad", cs.GetResult());
	//// set size
	//TheObjMeshManager::Instance()->Scale("title_quad", D3DXVECTOR3(TheScreenManager::Instance()->GetWindowWidth(),
	//	TheScreenManager::Instance()->GetWindowHeight(), 1.0f));
	//// draw
	////TheObjMeshManager::Instance()->SetColour("title_quad", D3DXVECTOR4(1.0f, 0.0f, 0.0f, 1.0f));
	//TheObjMeshManager::Instance()->Draw("title_quad");


	t.Translate(D3DXVECTOR3(0.0f, 0.0f, 0.0f));
	t.Draw();



	// draw the scene in 3d

	// draw scene in 3d end








	TheDX11ShaderManager::Instance()->ResetResources();
}

void TestState::Update()
{
	TheCameraManager::Instance()->Update();
	TheGame::Instance()->UpdateGameObjects();
}

void TestState::OnKeyDownEvent(const WPARAM& event)
{
	if (event == VK_RETURN)
	{
		m_retKeyDown = true;
	}
	else if (event == VK_ESCAPE)
	{
		TheGame::Instance()->SetGameState(TheOptionsState::Instance());
	}
}

void TestState::OnMouseMotionEvent(const LPARAM& event)
{
	int x = LOWORD(event);
	int y = HIWORD(event);
	
	int screenWidth = TheScreenManager::Instance()->GetWindowWidth();
	int screenHeight = TheScreenManager::Instance()->GetWindowHeight();
	
	if (x <= screenWidth && x >= 0)
	{
		m_mouseX = (float)x;
	}
	if (y <= screenHeight && y >= 0)
	{
		m_mouseY = (float)y;
	}
}
void TestState::OnMouseButtonDownEvent(const MSG& event)
{
	m_mouseDown = true;
}

void TestState::OnMouseButtonUpEvent(const MSG& event)
{
	m_mouseDown = false;
}

void TestState::OnActivated()
{
	TheEventPoller::Instance()->AddHandler(this);
}

void TestState::OnDeactivated()
{
	TheEventPoller::Instance()->RemoveHandler(this);
}

void TestState::GetAxisArray(Vec3f axes[])
{
	for (int i = 0; i != 15; i++)
	{
		axis[i] = axes[i];
	}
}
